#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Similarity calculation functions for item-based collaborative filtering recommendation.

Created on: 2025-04-20
Author: Nianqing Liu
"""

import numpy as np
from sklearn.metrics.pairwise import pairwise_distances
import scipy.sparse as sp


def cosine_similarity(ratings_matrix):
    """
    Calculate cosine similarity between items.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        Item-user rating matrix (transposed from user-item matrix).

    Returns:
    --------
    numpy.ndarray
        Item-item similarity matrix.
    """
    # Check if the matrix is sparse
    if sp.issparse(ratings_matrix):
        similarity = 1 - pairwise_distances(ratings_matrix, metric="cosine", n_jobs=-1)
    else:
        similarity = 1 - pairwise_distances(ratings_matrix, metric="cosine", n_jobs=-1)

    # Handle NaN values (if any)
    similarity = np.nan_to_num(similarity)

    return similarity


def pearson_similarity(ratings_matrix):
    """
    Calculate Pearson correlation between items.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        Item-user rating matrix (transposed from user-item matrix).

    Returns:
    --------
    numpy.ndarray
        Item-item similarity matrix.
    """
    # Calculate correlation coefficient
    similarity = np.corrcoef(ratings_matrix)

    # Handle NaN values (if any)
    similarity = np.nan_to_num(similarity)

    return similarity


def euclidean_similarity(ratings_matrix):
    """
    Calculate similarity based on Euclidean distance between items.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        Item-user rating matrix (transposed from user-item matrix).

    Returns:
    --------
    numpy.ndarray
        Item-item similarity matrix.
    """
    # Calculate Euclidean distance
    distances = pairwise_distances(ratings_matrix, metric="euclidean", n_jobs=-1)

    # Convert distance to similarity
    similarity = 1 / (1 + distances)

    return similarity


def adjusted_cosine_similarity(ratings_matrix, user_means=None):
    """
    Calculate adjusted cosine similarity between items.
    Normalizes ratings by subtracting user mean before calculation.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        Item-user rating matrix (transposed from user-item matrix).
    user_means : numpy.ndarray, optional
        Mean rating for each user. If None, calculates from ratings_matrix.

    Returns:
    --------
    numpy.ndarray
        Item-item similarity matrix.
    """
    # Transpose to get user-item matrix for calculating user means
    user_item_matrix = ratings_matrix.T

    # Get non-zero elements to calculate mean
    rated_mask = user_item_matrix != 0

    if user_means is None:
        # Calculate user means (only for rated items)
        user_sums = np.sum(user_item_matrix, axis=1)
        user_counts = np.sum(rated_mask, axis=1)
        # Avoid division by zero
        user_means = np.zeros_like(user_sums, dtype=float)
        nonzero_counts = user_counts > 0
        user_means[nonzero_counts] = (
            user_sums[nonzero_counts] / user_counts[nonzero_counts]
        )

    # Normalize user-item matrix by subtracting user mean
    normalized_matrix = np.zeros_like(user_item_matrix)
    for i, (user_ratings, mask, mean) in enumerate(
        zip(user_item_matrix, rated_mask, user_means)
    ):
        normalized_matrix[i, mask] = user_ratings[mask] - mean

    # Transpose back to get item-user normalized matrix
    normalized_item_user = normalized_matrix.T

    # Calculate cosine similarity on normalized item vectors
    similarity = 1 - pairwise_distances(
        normalized_item_user, metric="cosine", n_jobs=-1
    )

    # Handle NaN values (if any)
    similarity = np.nan_to_num(similarity)

    return similarity


def calculate_similarity(ratings_matrix, method="cosine", user_means=None):
    """
    Calculate item similarity matrix using the specified method.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    method : str
        Similarity method to use: 'cosine', 'pearson', 'euclidean', or 'adjusted_cosine'.
    user_means : numpy.ndarray, optional
        Mean rating for each user (used for adjusted_cosine).

    Returns:
    --------
    numpy.ndarray
        Item-item similarity matrix.
    """
    # Transpose the matrix to get item-user matrix
    item_user_matrix = ratings_matrix.T

    if method == "cosine":
        return cosine_similarity(item_user_matrix)
    elif method == "pearson":
        return pearson_similarity(item_user_matrix)
    elif method == "euclidean":
        return euclidean_similarity(item_user_matrix)
    elif method == "adjusted_cosine":
        return adjusted_cosine_similarity(item_user_matrix, user_means)
    else:
        raise ValueError(
            f"Unknown similarity method: {method}. Choose from 'cosine', 'pearson', 'euclidean', or 'adjusted_cosine'."
        )
